import os
import pickle
import torch
from deepctr_torch.models import DeepFM

with open("train_cached_minmax_分桶-1000.pkl", "rb")as f1:
    train_cached = pickle.load(f1)
with open("dev_test_cached_minmax_分桶-1000.pkl", "rb")as f1:
    dev_test_cached = pickle.load(f1)

DEVICE = "cuda:0" if torch.cuda.is_available() else 'cpu'


# 1. prelu-256-128
top10 = [5, 4, 1, 8, 22, 20, 29, 33, 15, 25]
label_weight_norm_manul = [0] + [0.2/10 if i in top10 else 0.8/(121-10) for i in range(1,122)]
label_weight = torch.tensor(label_weight_norm_manul).to(DEVICE)

model_path = "./save_deepFM_minmax-分桶-1000/"
if not os.path.exists(model_path): os.mkdir(model_path)
model_path += "DeepFM_prelu-256-128.pt"

model = DeepFM(linear_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_hidden_units=(256, 128), dnn_dropout=0.3, dnn_use_bn=True, dnn_activation='prelu',
                task='multiclass', device=DEVICE, class_num=122, model_save_path=model_path)
model.compile("adam", "cross_entropy", metrics=['acc_top1', 'acc_top3'])

model.fit(x=train_cached["train_model_input"], y=train_cached["df_train_target_values"],
        validation_data = (dev_test_cached["dev_model_input"], dev_test_cached["df_dev_target_values"]),
        batch_size=1024, epochs=20, verbose=1, label_weight=label_weight)


# test
for num_epoch in range(20):
    model_path = f"./save_deepFM_minmax-分桶-1000/DeepFM_prelu-256-128_epoch_{num_epoch}.pt"
    print(model_path)

    model = DeepFM(linear_feature_columns = dev_test_cached["fixlen_feature_columns"],
                    dnn_feature_columns = dev_test_cached["fixlen_feature_columns"],
                    dnn_hidden_units=(256, 128), dnn_dropout=0.3, dnn_use_bn=True, dnn_activation='prelu',
                    task='multiclass', device=DEVICE, class_num=122, model_save_path=model_path)
    model.compile("adam", "cross_entropy", metrics=['acc_top1', 'acc_top3'])

    model.load_state_dict((torch.load(model_path, map_location=torch.device(DEVICE))), strict=False)
    eval_result = model.evaluate(dev_test_cached["test_model_input"], dev_test_cached["df_test_target_values"],
                                 batch_size=1024)
    for k,v in eval_result.items():
        print(f"{k}: {v:.4f}")
    print("")
    
    
# 2. prelu-512-256
model_path = "./save_deepFM_minmax-分桶-1000/"
model_path += "DeepFM_prelu-512-256.pt"
model = DeepFM(linear_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_hidden_units=(512, 256), dnn_dropout=0.3, dnn_use_bn=True, dnn_activation='prelu',
                task='multiclass', device=DEVICE, class_num=122, model_save_path=model_path)
model.compile("adam", "cross_entropy", metrics=['acc_top1', 'acc_top3'])

model.fit(x=train_cached["train_model_input"], y=train_cached["df_train_target_values"],
        validation_data = (dev_test_cached["dev_model_input"], dev_test_cached["df_dev_target_values"]),
        batch_size=1024, epochs=20, verbose=1, label_weight=label_weight)
    
# test
for num_epoch in range(20):
    model_path = f"./save_deepFM_minmax-分桶-1000/DeepFM_prelu-512-256_epoch_{num_epoch}.pt"
    print(model_path)

    model = DeepFM(linear_feature_columns = dev_test_cached["fixlen_feature_columns"],
                    dnn_feature_columns = dev_test_cached["fixlen_feature_columns"],
                    dnn_hidden_units=(512, 256), dnn_dropout=0.3, dnn_use_bn=True, dnn_activation='prelu',
                    task='multiclass', device=DEVICE, class_num=122, model_save_path=model_path)
    model.compile("adam", "cross_entropy", metrics=['acc_top1', 'acc_top3'])

    model.load_state_dict((torch.load(model_path, map_location=torch.device(DEVICE))), strict=False)
    eval_result = model.evaluate(dev_test_cached["test_model_input"], dev_test_cached["df_test_target_values"],
                                 batch_size=1024)
    for k,v in eval_result.items():
        print(f"{k}: {v:.4f}")
    print("")

    
# 3. prelu-512-256-l2_reg-0.0001
model_path = "./save_deepFM_minmax-分桶-1000/"
model_path += "DeepFM_prelu-512-256-l2_reg-0.0001.pt"
model = DeepFM(linear_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_hidden_units=(512, 256), dnn_dropout=0.3, dnn_use_bn=True, dnn_activation='prelu',
                task='multiclass', device=DEVICE, class_num=122, model_save_path=model_path,
                l2_reg_linear=0.0001, l2_reg_embedding=0.0001)
model.compile("adam", "cross_entropy", metrics=['acc_top1', 'acc_top3'])

model.fit(x=train_cached["train_model_input"], y=train_cached["df_train_target_values"],
        validation_data = (dev_test_cached["dev_model_input"], dev_test_cached["df_dev_target_values"]),
        batch_size=1024, epochs=20, verbose=1, label_weight=label_weight)
    
# test
for num_epoch in range(20):
    model_path = f"./save_deepFM_minmax-分桶-1000/DeepFM_prelu-512-256-l2_reg-0.0001_epoch_{num_epoch}.pt"
    print(model_path)

    model = DeepFM(linear_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_feature_columns = train_cached["fixlen_feature_columns"],
                dnn_hidden_units=(512, 256), dnn_dropout=0.3, dnn_use_bn=True, dnn_activation='prelu',
                task='multiclass', device=DEVICE, class_num=122, model_save_path=model_path,
                l2_reg_linear=0.0001, l2_reg_embedding=0.0001)
    model.compile("adam", "cross_entropy", metrics=['acc_top1', 'acc_top3'])

    model.load_state_dict((torch.load(model_path, map_location=torch.device(DEVICE))), strict=False)
    eval_result = model.evaluate(dev_test_cached["test_model_input"], dev_test_cached["df_test_target_values"],
                                 batch_size=1024)
    for k,v in eval_result.items():
        print(f"{k}: {v:.4f}")
    print("")

    